from torch import nn


# 搭建神经网络
class CNN(nn.Module):
    def __init__(self, act_fun):  # 初始化函数
        super(CNN, self).__init__()  # super用于调用父类的初始化函数
        # 第一层
        self.layer1 = nn.Sequential(
            # 2维卷积，in_channel=1,out_channel=2
            nn.Conv2d(3, 18, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(18),
            eval(act_fun)()  # 激活函数
        )  # 3*6@28*28
        # 第二层
        self.layer2 = nn.Sequential(
            nn.MaxPool2d(kernel_size=2, stride=2),  # 池化（最大池化）
        )  # 3*6@14*14
        # 第三层
        self.layer3 = nn.Sequential(
            nn.Conv2d(18, 48, kernel_size=5, stride=1),  # 卷积
            nn.BatchNorm2d(48),  # 进行数据的归一化处理
            eval(act_fun)()
        )  # 16@10*10
        # 第四层
        self.layer4 = nn.Sequential(
            nn.MaxPool2d(kernel_size=2, stride=2)  # 池化
        )  # 3*16@5*5
        # 第五层（全连接层）
        self.fc = nn.Sequential(
            nn.Linear(228528, 120),
            eval(act_fun)(),
            nn.Linear(120, 84),
            nn.Sigmoid(),            # nn.Softmax(),
            nn.Linear(84, 2)
        )

    def forward(self, x):  # 传播
        x = self.layer1(x)  # 6@28*28
        x = self.layer2(x)  # 6@14*14
        x = self.layer3(x)  # 16@10*10
        x = self.layer4(x)  # 16@5*5
        x = x.view(x.size(0), -1)  # 将多维度的tensor展平成一维
        x = self.fc(x)
        return x
